package nachos.vm;

import nachos.machine.*;
import nachos.threads.*;
import nachos.userprog.*;

import java.util.*;

/**
 * A kernel that can support multiple demand-paging user processes.
 */
public class VMKernel extends UserKernel
{
    /**
     * Allocate a new VM kernel.
     */
    public VMKernel()
    {
        super();

        for (int i = 0; i < coreMap.length; i++) {
            coreMap[i] = new MemoryEntry(i);
        }
    }

    /**
     * Initialize this kernel.
     */
    public void initialize(String[] args)
    {
        super.initialize(args);
        memLock = new Lock();
        allPinned = new Condition(memLock);
        swap = new Swap();
    }

    /**
     * Test this kernel.
     */
    public void selfTest()
    {
        super.selfTest();
    }

    /**
     * Start running user programs.
     */
    public void run()
    {
        super.run();
    }

    /**
     * Terminate this kernel. Never returns.
     */
    public void terminate()
    {
        swap.cleanup();

        super.terminate();
    }

    private MemoryEntry clockAlgorithm() {
        memLock.acquire();
        while (pinnedCount == coreMap.length) {
            // 等待有空位
            allPinned.sleep();
        }

        propagateAndFlushTLB(false);

        MemoryEntry entry;

        // 找到一个未被锁的页
        while (true) {
            idx = (idx + 1) % coreMap.length;
            entry = coreMap[idx];

            if (entry.pinned) {
                continue;
            }

            // 失效的页面可以使用
            if (entry.pid == -1 || !entry.translationEntry.valid) {
                break;
            }

            // 最近使用过，可能会再次被使用
            if (entry.translationEntry.used) {
                entry.translationEntry.used = false;
            } else {
                break;
            }
        }

        // 锁住该页
        entry.pinned = true;
        pinnedCount++;

        invalidateTLBEntry(idx);

        MemoryEntry entry1 = null;
        if (entry.pid > -1) {
            entry1 = invertedPageTable.remove(new TableKey(entry.translationEntry.vpn, entry.pid));
        }

        memLock.release();

        if (entry1 != null) {
            swap.swapOut(entry);
        }

        return entry;
    }

    void propagateAndFlushTLB(boolean flush) {
        for (int i = 0; i < Machine.processor().getTLBSize(); i++) {
            TranslationEntry memEntry = Machine.processor().readTLBEntry(i);

            if (memEntry.valid) {
                TranslationEntry mapEntry = coreMap[memEntry.ppn].translationEntry;
                if (mapEntry.valid && mapEntry.vpn == memEntry.vpn) {
                    mapEntry.used = memEntry.used | mapEntry.used;
                    mapEntry.dirty = memEntry.dirty | mapEntry.dirty;
                }
            }

            if (flush) {
                memEntry.valid = false;
                Machine.processor().writeTLBEntry(i, memEntry);
            }
        }
    }

    void restoreTLB() {
        for (int i = 0; i < Machine.processor().getTLBSize(); i++) {
            TranslationEntry entry = coreMap[i].translationEntry;
            Machine.processor().writeTLBEntry(i, entry);
        }
    }

    void invalidateTLBEntry(int ppn) {
        for (int i = 0; i < Machine.processor().getTLBSize(); i++) {
            TranslationEntry translationEntry = Machine.processor().readTLBEntry(i);
            if (translationEntry.ppn == ppn && translationEntry.valid) {
                translationEntry.valid = false;
                Machine.processor().writeTLBEntry(i, translationEntry);
                break;
            }
        }
    }

    TranslationEntry requestFreePage(int vpn, int pid) {
        MemoryEntry page = clockAlgorithm();

        // 初始化这块内存区域
        int pagePos = Processor.makeAddress(page.translationEntry.ppn, 0);
        Arrays.fill(Machine.processor().getMemory(), pagePos, pagePos + Processor.pageSize, (byte) 0);

        page.translationEntry.vpn = vpn;
        page.translationEntry.valid = true;
        page.pid = pid;

        // 加进页表
        insertIntoTable(vpn, pid, page);

        return page.translationEntry;
    }

    private void insertIntoTable(int vpn, int pid, MemoryEntry page) {
        memLock.acquire();
        invertedPageTable.put(new TableKey(vpn, pid), page);
        memLock.release();
    }

    TranslationEntry pageFault(int vpn, int pid) {
        if (!swap.pageInSwap(vpn, pid)) {
            return null;
        }

        TranslationEntry entry = requestFreePage(vpn, pid);
        swap.swapIn(vpn, pid, entry.ppn);

        return entry;
    }

    void freePages(int pid, int maxVPN) {
        memLock.acquire();
        for (int i = 0; i < coreMap.length; i++) {
            MemoryEntry page = coreMap[i];
            if (page.pid == pid) {
                invertedPageTable.remove(new TableKey(page.translationEntry.vpn, page.pid));

                page.pid = -1;
                page.translationEntry.valid = false;
            }
        }
        memLock.release();

        swap.freePages(maxVPN, pid);
    }

    void unpin(int ppn) {
        memLock.acquire();
        MemoryEntry entry = coreMap[ppn];

        if (entry.pinned) {
            entry.pinned = false;
            pinnedCount--;
        }

        allPinned.wake();

        memLock.release();
    }

    TranslationEntry pinIfExists(int vpn, int pid) {
        MemoryEntry entry;

        memLock.acquire();
        if ((entry = invertedPageTable.get(new TableKey(vpn, pid))) != null) {
            if (!entry.pinned) {
                entry.pinned = true;
                pinnedCount++;
            }
        }
        memLock.release();

        if (entry == null) {
            return null;
        } else {
            return entry.translationEntry;
        }
    }

    void propagateEntry(int ppn, boolean used, boolean dirty) {
        memLock.acquire();
        TranslationEntry entry = coreMap[ppn].translationEntry;
        entry.used = used | entry.used;
        entry.dirty = dirty | entry.dirty;
        memLock.release();
    }


    // dummy variables to make javac smarter
    private static VMProcess dummy1 = null;

    private static final char dbgVM = 'v';

    private MemoryEntry[] coreMap = new MemoryEntry[Machine.processor().getNumPhysPages()];

    private int idx = 0;

    // 进程的虚拟页号组成的TableKey与物理页号映射
    // 全局共享
    private Map<TableKey, MemoryEntry> invertedPageTable = new Hashtable<>();

    private Lock memLock;

    private int pinnedCount;

    private Condition allPinned;

    private Swap swap;

    // --------------------------------------------
    // Inner classes
    // --------------------------------------------

    private static class TableKey {
        TableKey(int vpn, int pid) {
            this.vpn = vpn;
            this.pid = pid;
        }

        @Override
        public int hashCode() {
            return Processor.makeAddress(vpn, pid);
        }

        @Override
        public boolean equals(Object x) {
            if (this == x)
                return true;
            else if (x instanceof TableKey) {
                TableKey xCasted = (TableKey)x;
                return vpn == xCasted.vpn && pid == xCasted.pid;
            } else {
                return false;
            }
        }

        private int vpn, pid;
    }

    private class Swap {
        Swap() {
            swapFile = fileSystem.open("swapfile", true);
        }

        void swapOut(MemoryEntry entry) {
            if (entry.translationEntry.valid) {

                SwapEntry swapEntry = null;
                TableKey key = new TableKey(entry.translationEntry.vpn, entry.pid);

                swapLock.acquire();
                if (entry.translationEntry.dirty || !swapTable.containsKey(key)) {
                    if (freeList.size() > 0) {
                        swapEntry = freeList.removeFirst();
                        swapEntry.readOnly = entry.translationEntry.readOnly;
                    } else {
                        swapEntry = new SwapEntry(maxTableEntry++, entry.translationEntry.readOnly);
                    }

                    swapTable.put(key, swapEntry);
                }
                swapLock.release();

                if (swapEntry != null) {
                    Lib.assertTrue(
                            swapFile.write(swapEntry.swapPageNumber * Processor.pageSize,
                                    Machine.processor().getMemory(),
                                    entry.translationEntry.ppn * Processor.pageSize,
                                    Processor.pageSize
                                    ) == Processor.pageSize
                    );
                }
            }
        }


        void swapIn(int pid, int vpn, int ppn) {
            swapLock.acquire();
            SwapEntry swapEntry = swapTable.get(new TableKey(vpn, pid));
            swapLock.release();

            if (swapEntry != null) {
                // 保证从虚拟内存中读出的数据读满一个页面大小
                Lib.assertTrue(swapFile.read(
                        swapEntry.swapPageNumber * Processor.pageSize,
                        Machine.processor().getMemory(),
                        ppn * Processor.pageSize,
                        Processor.pageSize
                ) == Processor.pageSize);

                coreMap[ppn].translationEntry.readOnly = swapEntry.readOnly;
            }
        }

        boolean pageInSwap(int vpn, int pid) {
            swapLock.acquire();
            boolean inSwap = swapTable.containsKey(new TableKey(vpn, pid));
            swapLock.release();

            return inSwap;
        }

        void freePages(int maxVpn, int pid) {
            swapLock.acquire();
            SwapEntry entry;
            for (int i = 0; i < maxVpn; i++) {
                if ((entry = swapTable.get(new TableKey(i, pid))) != null) {
                    freeList.add(entry);
                }
            }
            swapLock.release();
        }

        void cleanup() {
            swapFile.close();
            fileSystem.remove(swapFile.getName());
        }

        private OpenFile swapFile;

        // 当前交换文件中可用的地址
        private LinkedList<SwapEntry> freeList = new LinkedList<>();

        // 进程页号和虚拟内存中地址的映射表
        private HashMap<TableKey, SwapEntry> swapTable = new HashMap<>();

        // 保证对swapTable的原子性操作
        private Lock swapLock = new Lock();
        private int maxTableEntry = 0;

        private class SwapEntry {
            SwapEntry(int swapPageNumber, boolean readOnly) {
                this.swapPageNumber = swapPageNumber;
                this.readOnly = readOnly;
            }

            int swapPageNumber;
            boolean readOnly;
        }
    }

    private static class MemoryEntry {
        MemoryEntry(int ppn) {
            translationEntry = new TranslationEntry(-1, ppn, false, false, false, false);
        }

        TranslationEntry translationEntry;
        int pid = -1;
        boolean pinned = false;
    }
}
